from manimlib import *
import math

__curr_dir = os.path.dirname(__file__)
if __curr_dir not in sys.path:
    sys.path.append(__curr_dir)


# 弹簧圈数
turns = 20


def Spring3D(start, end, axis):
    """
    创建弹簧，包括首尾节点（球体）和长度标签
    """

    start = np.array(start)
    end = np.array(end)

    global_length = np.linalg.norm(end - start)

    start = axis.c2p(start[0], start[1], start[2])
    end = axis.c2p(end[0], end[1], end[2])

    # 横截面半径
    radius = global_length * 0.05

    curve = mobject.functions.ParametricCurve(
        lambda t: axis.c2p(
            t, radius * np.cos(2 * np.pi * turns * t), radius * np.sin(2 * np.pi * turns * t)  # x坐标  # y坐标  # z坐标
        ),
        t_range=[0, 1, 0.005],  # 绘制步长0.001
        color=YELLOW,
    )
    curve.set_style(stroke_width=10)

    # 将单元弹簧映射到真实坐标上
    curve.shift(start)  # 将弹簧移动到起始位置
    curve.stretch(global_length, 0, about_point=start)  # 根据向量的长度来缩放弹簧

    # # 在弹簧的两端添加两个球体
    # ball_start = Dot(radius=radius * 0.2, color=RED).move_to(start)

    return curve


class TwoDimensionSpring(Scene):
    def construct(self):
        axes = Axes(
            x_range=(-0.2, 1.2, 0.2),
            y_range=(-0.2, 1.2, 0.2),
            axis_config=dict(
                stroke_color=GREY_A,
                numbers_to_exclude=[0],
                tick_size=0.02,
            ),
        )
        frame = self.camera.frame
        frame.set_height(2)
        frame.move_to(axes.coords_to_point(1, 0.7))

        # ========================================================
        # ========================================================
        # ========================================================
        # 创建弹簧
        sprint_start = [0, 0, 0]
        sprint_end = [1, 0, 0]
        spring = Spring3D(sprint_start, sprint_end, axes)
        spring.move_to(axes.c2p(0.5, 0, 0))
        spring_1d_k = Tex(
            r"k\begin{bmatrix} 1.00 & -1.00 \\  -1.00  &  1.00  \end{bmatrix}",
            font_size=10,
        )
        spring_1d_k.move_to(axes.c2p(1.5, 1, 0))

        spring_2d_k = Tex(
            r"""k\begin{bmatrix}
            1.00 & 0.00 & -1.00 & 0.00 \\
            0.00 & 0.00 & 0.00 & 0.00 \\
            -1.00 & 0.00 & 1.00 & 0.00 \\
            0.00 & 0.00 & 0.00 & 0.00 \end{bmatrix}""",
            font_size=10,
        )
        spring_2d_k.move_to(axes.c2p(1.5, 1, 0))

        # 1/4圆
        arc = axes.get_graph(lambda x: (1 - x**2) ** 0.5, x_range=[0, 1, 0.01])

        # 加一个点，用来计算theta
        dot = Dot(fill_color=RED, radius=0.01, stroke_width=0, fill_opacity=0.2)
        dot.move_to(axes.i2gp(1, arc))
        x_tracker = ValueTracker(1)
        f_always(dot.move_to, lambda: axes.i2gp(x_tracker.get_value(), arc))

        self.add(axes[0], spring, spring_1d_k)
        self.wait(2)

        theta_tex = Tex(r"\theta=  0\enspace ^\circ", font_size=10)
        theta_value = theta_tex.make_number_changable("0")
        theta_tex.move_to(axes.coords_to_point(1, 1.4))

        self.play(
            ShowCreation(axes[1]),
            ShowCreation(arc),
            ShowCreation(theta_tex),
            FadeIn(dot, scale=0.5),
            ReplacementTransform(spring_1d_k, spring_2d_k),
        )
        self.wait(2)

        def get_theta():
            coord = dot.get_center() - axes.get_origin()
            x = coord[0]
            y = coord[1]
            return math.atan2(y, x)

        def get_theta_degree():
            # 将弧度转换为度
            return math.degrees(get_theta())

        def coscos():
            return (math.cos(get_theta())) ** 2

        theta_value.f_always.set_value(get_theta_degree)

        # coscos_nums = spring_2d_k.make_number_changable("0.00", replace_all=True)
        # for coscos_num in coscos_nums:
        #     coscos_num.f_always.set_value(coscos)

        self.play(
            Rotate(spring, np.pi / 4, about_point=axes.coords_to_point(0, 0, 0)),
            x_tracker.animate.set_value((0.5) ** 0.5),
            run_time=2,
        )
        self.wait(2)
        self.play(
            Rotate(spring, np.pi / 4, about_point=axes.coords_to_point(0, 0, 0)),
            x_tracker.animate.set_value(0),
            run_time=2,
        )
        self.wait(2)
        self.play(
            Rotate(spring, -np.pi / 4, about_point=axes.coords_to_point(0, 0, 0)),
            x_tracker.animate.set_value((0.5) ** 0.5),
            run_time=2,
        )
        self.wait(2)
        self.play(
            Rotate(spring, -np.pi / 4, about_point=axes.coords_to_point(0, 0, 0)),
            x_tracker.animate.set_value(1),
            run_time=2,
        )
        self.wait(2)
